import argparse
import os
from enum import Enum
from time import time
from itertools import chain


# What to find

# name: list of name variants
# pins: list of needed pins
# Example: BlockToFind(name=["UART"], inouts=["RX", "TX"])
# There may be no inouts required, ex. timer for DMA req generation only:
#   BlockToFind(name="TIMER", dmas=["CH0", "CH1", "CH2", "CH3"])
# Signals may be omitted if same as inouts. Needed for ex. timer with 4 channels and 2 outputs:
#   BlockToFind(name="TIMER", signals=["CH0", "CH1", "CH2", "CH3"], inouts=["CH0", "CH2"])
# There may be several variants of resources required, generally useful for timers:
#   BlockToFind("TIMER", [PinSigDma(pins=["CH0"]), PinSigDma(pins=["CH1"]), PinSigDma(pins=["CH2"])]),

class PinSigDma:
    def __init__(self, pins=None, signals=None, dmas=None):
        self.pins: list[str] = [] if pins is None else pins
        self.signals: list[str] = pins if signals is None else signals
        self.dmas: list[str] = [] if dmas is None else dmas


# Name_str examoles: "UART*" "DAC*" "TIMER*" "TIMER0, TIMER1, TIMER2". UART means both UART and USART.
class BlockToFind:
    def __init__(self, name_str: str, pin_sig_dmas: list[PinSigDma]):
        self.name_str = name_str
        self.pin_sig_dmas: list[PinSigDma] = pin_sig_dmas
        self.names: list[str] = []  # Will be filled later


# ADC, DAC, I2C, I2S, SDIO, SPI, TIMER, UART
# All save timers
needed_blocks_without_timers = [
    BlockToFind("UART*", [PinSigDma(pins=["RX", "TX"], dmas=["RX", "TX"])]),
    BlockToFind("DAC", [  # IR TX
        PinSigDma(pins=["OUT0"], dmas=["CH0"]),
        PinSigDma(pins=["OUT1"], dmas=["CH1"])
    ]),
    BlockToFind("I2S*",  [PinSigDma(pins=["WS", "CK", "SD"], dmas=["TX"])]),
    BlockToFind("SDIO", [PinSigDma(pins=["D0", "D1", "D2", "D3", "CK", "CMD"], dmas=["SDIO"])]),

    # WS2812
    BlockToFind("SPI*",  [PinSigDma(pins=["MOSI"], dmas=["TX"])]),  # External
    BlockToFind("SPI*",  [PinSigDma(pins=["MOSI"], dmas=["TX"])]),  # Side
]

needed_timers = [
    # Side and Front: Timer must have DMACFG reg
    # BlockToFind("TIMER0, TIMER7, TIMER1, TIMER2, TIMER3, TIMER4", [
    #     PinSigDma(pins=["CH0", "CH1"], dmas=["UP"]),  # Only sequental channels are suitable,
    #     PinSigDma(pins=["CH1", "CH2"], dmas=["UP"]),  # as DMA burst will fill adjacent registers.
    #     PinSigDma(pins=["CH2", "CH3"], dmas=["UP"]),
    # ]),
    BlockToFind("TIMER*", [  # Front
        PinSigDma(pins=["CH0"], dmas=["UP"]),
        PinSigDma(pins=["CH1"], dmas=["UP"]),
        PinSigDma(pins=["CH2"], dmas=["UP"]),
        PinSigDma(pins=["CH3"], dmas=["UP"]),
    ]),

    # IR RX x4
    BlockToFind("TIMER*", [
        PinSigDma(pins=["CH0"], signals=["CH0", "CH1"]),
        PinSigDma(pins=["CH1"], signals=["CH0", "CH1"]),
    ]),
    BlockToFind("TIMER*", [
        PinSigDma(pins=["CH0"], signals=["CH0", "CH1"]),
        PinSigDma(pins=["CH1"], signals=["CH0", "CH1"]),
    ]),
    BlockToFind("TIMER*", [
        PinSigDma(pins=["CH0"], signals=["CH0", "CH1"]),
        PinSigDma(pins=["CH1"], signals=["CH0", "CH1"]),
    ]),
    BlockToFind("TIMER*", [
        PinSigDma(pins=["CH0"], signals=["CH0", "CH1"]),
        PinSigDma(pins=["CH1"], signals=["CH0", "CH1"]),
    ]),

    # PWM Input
    BlockToFind("TIMER*", [
        PinSigDma(pins=["CH0"], signals=["CH0", "CH1"], dmas=["CH0"]),
        PinSigDma(pins=["CH1"], signals=["CH0", "CH1"], dmas=["CH1"]),
    ]),

    # Sys timer
    BlockToFind("TIMER*", [
        PinSigDma(signals=["CH0"]),
        PinSigDma(signals=["CH1"]),
        PinSigDma(signals=["CH2"]),
        PinSigDma(signals=["CH3"]),
    ]),
]

gpios_in_use = [
    "PA11", "PA12",  # USB
    "PA13", "PA14",  # SWD
    "PD0", "PD1",    # Crystal
]

# Do not touch, will be filled automatically
pinnums_in_use: list[int] = []


# region === Raw data lists ====
class PinInfo:
    number: int = -1
    name = ""
    kind = ""
    default = ""
    is_5v_tol = "false"

    def __init__(self):
        self.alternate = []
        self.remap = []

    def GetGpio(self):
        def CheckIfGpio(s: str):
            return s[0] == 'P' and s[1] in ['A', 'B', 'C', 'D', 'E', 'F', 'G']
        if CheckIfGpio(self.name):
            return self.name
        else:  # Name didn't work, try remapped
            for rmp in self.remap:
                if CheckIfGpio(rmp):
                    return rmp
        return None

    def Print(self):
        print("{} {} {} {}; ".format(self.number, self.name, self.kind, self.default), end="")
        if len(self.alternate) > 0:
            print("Alt: ", end="")
            for s in self.alternate:
                print("{}; ".format(s), end="")
        if len(self.remap) > 0:
            print("Rmp: ", end="")
            for s in self.remap:
                print("{}; ".format(s), end="")
        print("\n", end="")


class DMAInfo:
    def __init__(self, per_name, per_chnl, dma_n, dma_chnl):
        self.periph_name = per_name
        self.periph_chnl = per_chnl
        self.dma_n = dma_n
        self.dma_chnl = dma_chnl


class GpioSignal:
    def __init__(self, signal, gpio):
        self.signal = signal
        self.gpio = gpio


class RemapInfo:
    def __init__(self, name, option, pins):
        self.periph_name = name
        self.option = option
        self.signals: list[GpioSignal] = []  # Signal name, GPIO
        if len(pins) == 1:
            self.signals.append(GpioSignal("CH0", pins[0]))  # Signal name, GPIO
        else:
            for pin in pins:
                tokens = pin.strip(" ,").split("/")
                gpio = tokens[1]
                # Several signals possible: TIMER1_CH0-TIMER1_ETI PA0
                sigs = tokens[0].split('-')  # TIMER1_CH0 TIMER1_ETI
                for sig in sigs:
                    sig_name = sig.split("_", 1)[1]
                    self.signals.append(GpioSignal(sig_name, gpio))


pin_list: list[PinInfo] = []
dma_list: list[DMAInfo] = []
remap_list: list[RemapInfo] = []
all_gpios: list[str] = []
# endregion


def GpioName2PinNum(gpio_name: str):
    for pin in pin_list:
        if pin.GetGpio() == gpio_name:
            return pin.number
    return None


class BlockDma:
    def __init__(self, sig_name, dma_num, dma_ch):
        self.signal_name = sig_name
        self.dma_n = dma_num
        self.dma_chnl = dma_ch

    def ToString(self):
        return "{}_{}ch{}".format(self.signal_name, self.dma_n, self.dma_chnl)

    def ToDmaNAndChnl(self):
        return "{}{}".format(self.dma_n, self.dma_chnl)

    def ToInt(self):
        return int(self.dma_n) * 100 + int(self.dma_chnl)


class ConfiguredBlock:
    def __init__(self, name, indx, remap_option: int, gpios: list[str], dmas: list[BlockDma]):
        self.name = name
        self.remap_option = remap_option
        self.dmas: list[BlockDma] = dmas.copy()
        self.gpios: list[str] = gpios.copy()
        # Ints
        self.block_indx = indx
        self.pinnums = [int(GpioName2PinNum(gpio_name)) for gpio_name in gpios]
        self.int_dmas = [int(dma.ToInt()) for dma in dmas]


class GpiosVariant:
    def __init__(self, gpios_signals: list[GpioSignal], remap_option: int):
        self.gpios_signals: list[GpioSignal] = gpios_signals
        self.remap_option: int = remap_option


class PeriphBlock:
    def __init__(self, name, indx):
        self.name: str = name  # USART1, TIMER2, SDIO...
        self.gpios_variants: list[GpiosVariant] = []  # Several variants of signals connection
        self.dmas: list[BlockDma] = []  # list of DMA signals usable
        self.block_indx: int = indx

    def GetVariantsForPinsAndDmas(self, pin_sig_dma: PinSigDma):
        # Do we have all needed dmas available simultaneously? I.e. Timer's CH2 and UP at same channel are not ok.
        dmas: list[BlockDma] = []
        for dma_name in pin_sig_dma.dmas:
            for dma in self.dmas:
                if dma.signal_name == dma_name:
                    # Check if channel is already occupied
                    for dma_b in dmas:
                        if dma != dma_b and dma.dma_n == dma_b.dma_n and dma.dma_chnl == dma_b.dma_chnl:
                            return []  # Occupied DMA channel found
                    dmas.append(dma)  # not duplicated
                    break
            else:  # Needed DMA not found
                return []

        var_list: list[ConfiguredBlock] = []
        # Add all variants having needed signals and pins
        for pin_var in self.gpios_variants:
            # Check if all requested signals present in this variant
            for signal in pin_sig_dma.signals:
                for gpio_sig in pin_var.gpios_signals:
                    if gpio_sig.signal == signal:
                        break
                else:  # Not found
                    break  # no need to check other signals if one is not found
            else:  # all signals were found, add requested pins
                gpios: list[str] = []
                for pin in pin_sig_dma.pins:
                    for gpio_sig in pin_var.gpios_signals:
                        if gpio_sig.signal == pin and gpio_sig.gpio in all_gpios:
                            gpios.append(gpio_sig.gpio)
                            break  # Signal found
                    else:  # Not found
                        gpios.clear()
                        break  # Stop iterating pins
                else:  # All pins were iterated and found
                    # Check if gpios are distinct: i.e. Timer's CH0 and ETI at same GPIO is not ok.
                    if len(gpios) == len(set(gpios)):  # no duplicates, good variant
                        var_list.append(ConfiguredBlock(self.name, self.block_indx, pin_var.remap_option, gpios, dmas))
        return var_list

    def Print(self):
        print(self.name, " ", end="")
        # Pins
        print("[", end="")
        for gpios_var in self.gpios_variants:
            print("[{}".format(gpios_var.remap_option), end="")
            for gpio_sig in gpios_var.gpios_signals:
                print("({} {})".format(gpio_sig.signal, gpio_sig.gpio), end="")
            print("] ", end="")
        print("]")
        # DMAs
        # print("[", end="")
        # for dma in self.dmas:
        #     print("<", dma.signal_name, dma.dma_n, dma.dma_chnl, "> ", end="")
        # print("]")


all_blocks: list[PeriphBlock] = []


def ConstructAllBlocksAvailable():
    # region Find block names
    block_names = set()

    def AddName(asignal: str):
        ftokens = asignal.split('_')  # TIMER2_CH1, ADC01_IN8, SDIO_D2, PC11
        if len(ftokens) != 2:  # PC11
            return
        # Special cases
        match(ftokens[0]):
            case "ADC012":
                block_names.add("ADC0")
                block_names.add("ADC1")
                block_names.add("ADC2")
            case "ADC01":
                block_names.add("ADC0")
                block_names.add("ADC1")
            case _:
                block_names.add(ftokens[0])

    for pin in pin_list:
        for signal in pin.alternate:
            AddName(signal)
        for signal in pin.remap:
            AddName(signal)

    # Sort it
    block_names = list(block_names)
    block_names.sort()
    # Put to blocks
    for indx, name in enumerate(block_names):
        all_blocks.append(PeriphBlock(name, indx))
    # endregion

    def NameIsEqual(block_name, signal_name):
        return (block_name == signal_name or
                (block_name == "ADC0" and signal_name in ["ADC01", "ADC012"]) or
                (block_name == "ADC1" and signal_name in ["ADC01", "ADC012"]) or
                (block_name == "ADC2" and signal_name == "ADC012"))

    # region Add pins to block
    for block in all_blocks:
        # Remap_data includes all pinning variants, so start with remap
        found = False
        for rmp in remap_list:
            if NameIsEqual(block.name, rmp.periph_name):
                found = True
                block.gpios_variants.append(GpiosVariant(rmp.signals, rmp.option))
        if not found:  # Now try alternative function
            signals: list[GpioSignal] = []
            for pin in pin_list:
                for alt in pin.alternate:
                    tokens = alt.split('_', 1)
                    if NameIsEqual(block.name, tokens[0]):
                        signals.append(GpioSignal(tokens[1], pin.GetGpio()))
            if len(signals) > 0:
                block.gpios_variants.append(GpiosVariant(signals, 0))
    # endregion
    # region Add DMAs to block
    for block in all_blocks:
        for dma in dma_list:
            if dma.periph_name == block.name:
                block.dmas.append(BlockDma(sig_name=dma.periph_chnl, dma_num=dma.dma_n, dma_ch=dma.dma_chnl))
    # endregion


def ReadPins(lines):
    class State(Enum):
        kName = 0
        kNumber = 1
        kKind = 2
        kOther = 3

    pin = PinInfo()
    sta = State.kName
    for line in lines:
        tokens = line.strip().split()
        if len(tokens) == 0:
            continue
        # Some regions
        if tokens[0] == "===Pins===":
            continue  # Just ignore
        elif tokens[0] == "===DMA===":
            break
        # Ordinal situation, proceed parsing
        match sta:
            case State.kName:
                pin.name = tokens[0].split('-')[0]  # "PC13-TAMPER-RTC" -> "PC13"
                sta = State.kNumber
            case State.kNumber:
                pin.number = tokens[0]
                sta = State.kKind
            case State.kKind:
                pin.kind = tokens[0]
                sta = State.kOther
            case State.kOther:  # Default, alternate, remap
                if tokens[0].lower().startswith("default"):
                    pin.default = tokens[1].strip(',')  # "Default: VBAT"
                elif tokens[0].lower().startswith("alternate"):
                    for s in tokens[1:]:
                        pin.alternate.append(s.strip(',').split('(')[0])
                elif tokens[0].lower().startswith("remap"):
                    for s in tokens[1:]:
                        pin.remap.append(s.strip(',').split('(')[0])  # Remove ',' and things like (5)
                elif tokens[0].lower() == "5vt":
                    pin.is_5v_tol = "true"
                elif tokens[0] == '-':
                    continue  # Just ignore
                else:  # New name
                    pin_list.append(pin)
                    pin = PinInfo()
                    pin.name = tokens[0].split('-')[0]  # "PC13-TAMPER-RTC" -> "PC13"
                    sta = State.kNumber
    pin_list.append(pin)  # Add last pin


def ReadDMA(lines):
    dma_n = -1
    for line in lines:
        line = line.strip()
        # Process region start or DMA region end
        if line == "DMA0":
            dma_n = 0
        elif line == "DMA1":
            dma_n = 1
        elif line == "===REMAP===":
            return
        elif dma_n != -1:
            # We inside DMA0 or DMA1 region
            tokens = line.split()  # TIMER0_CH0 1
            if len(tokens) < 2:
                continue
            # TIMER0_CH0 or ADC0?
            tokens2 = tokens[0].split('_')
            if len(tokens2) == 1:  # ADC0, ADC2, SDIO etc.
                dma_list.append(DMAInfo(per_name=tokens[0], per_chnl=tokens[0], dma_n=dma_n, dma_chnl=int(tokens[1])))
            else:  # TIMER0_CH0, USART1_RX etc.
                dma_list.append(DMAInfo(per_name=tokens2[0], per_chnl=tokens2[1], dma_n=dma_n, dma_chnl=int(tokens[1])))


def ReadRemap(lines):
    found = False
    for line in lines:
        line = line.strip()
        # Find Remap region start
        if line == "===REMAP===":
            found = True
            continue
        elif not found:
            continue
        # We inside REMAP region
        tokens = line.split()
        if len(tokens) < 3:
            continue
        # Some tokens found
        remap_list.append(RemapInfo(name=tokens[0], option=int(tokens[1]), pins=tokens[2:]))


def FindCombination(list_of_block_sets: list[list[ConfiguredBlock]]):
    indx_list: list[int] = [0] * len(list_of_block_sets)
    comb_set = []
    n = 0

    def CheckNames():
        indxs = []
        for blist, indx in zip(list_of_block_sets, indx_list):
            indxs.append(blist[indx].block_indx)
        return len(indxs) == len(set(indxs))  # Names are distinct

    def CheckDmas():
        idmas = []
        for blist, indx in zip(list_of_block_sets, indx_list):
            idmas.extend(blist[indx].int_dmas)
        return len(idmas) == len(set(idmas))

    def CheckGpios():
        pinnums = pinnums_in_use.copy()
        for blist, indx in zip(list_of_block_sets, indx_list):
            pinnums.extend(blist[indx].pinnums)
        return len(pinnums) == len(set(pinnums))

    def PrintAndAddComb():
        print("\n{}: ".format(n), end="")
        iblock_set = []
        for set_n, block_indx in enumerate(indx_list):
            iblock = list_of_block_sets[set_n][block_indx]
            iblock_set.append(iblock)
            dma_str = " ".join([dma.ToString() for dma in iblock.dmas])
            gpio_str = " ".join(iblock.gpios)
            if len(dma_str) > 0:
                print("{}_r{} dma[{}] g[{}]; ".format(iblock.name, iblock.remap_option, dma_str, gpio_str), end="")
            else:
                print("{}_r{} g[{}]; ".format(iblock.name, iblock.remap_option, gpio_str), end="")
        comb_set.append(iblock_set)

    checked = 0
    start_time = time()
    while True:
        checked += 1
        if CheckNames():
            if CheckDmas():
                if CheckGpios():
                    n += 1
                    PrintAndAddComb()

        # Construct new indxes: 0 0 0 -> 1 0 0 -> 2 0 0 -> 0 1 0 -> 1 1 0 -> 2 1 0 -> 0 0 1 -> 1 0 1...
        for i in range(len(indx_list)):
            indx_list[i] += 1
            if indx_list[i] < len(list_of_block_sets[i]):
                break
            indx_list[i] = 0
        else:
            break

    print("\n\nTime consumed: {}".format(time() - start_time))
    print("Combinations checked: {}".format(checked))
    if n == 0:
        print("No combinations found")
    return comb_set


def main():
    print(".txt file to pins file")
    parser = argparse.ArgumentParser(description=".txt file to pins json file")
    parser.add_argument("-t", "--txt", default="", help='.txt file')
    args = parser.parse_args()

    # region === Check args ===
    if args.txt == "":
        print("Specify .txt file")
        return
    if args.txt and (not os.path.exists(args.txt) or not os.path.isfile(args.txt)):
        print("Bad .txt file")
        return
    # endregion

    # region === Read Data ===
    with open(args.txt, "r") as f:
        lines = f.readlines()

    # Find MCU
    mcu = ""
    tokens = lines[0].strip().split()
    if len(tokens) >= 2 and tokens[0] == "MCU" and tokens[1] != "":
        mcu = tokens[1]
        lines = lines[1:]
    else:
        return
    print(mcu)

    ReadPins(lines)
    ReadDMA(lines)
    ReadRemap(lines)

    # Fill GPIOs available
    for pin in pin_list:
        s = pin.GetGpio()
        if s:
            all_gpios.append(s)
    all_gpios.sort()

    ConstructAllBlocksAvailable()
    # endregion

    # Fill pinnums_in-use
    for gpio_name in gpios_in_use:
        for pin in pin_list:
            if pin.GetGpio() == gpio_name:
                pinnums_in_use.append(int(pin.number))

    # Construct block names out of name_str
    def GetBlockNamesStartingWith(aname: str):
        if aname.endswith('*'):  # TIMER*  : If name ends with '*', find all suitable names
            return [ablock.name for ablock in all_blocks if ablock.name.startswith(aname[:-1])]
        else:  # TIMER1
            return [ablock.name for ablock in all_blocks if ablock.name == aname]

    for needed_block in chain(needed_blocks_without_timers, needed_timers):
        tokens = needed_block.name_str.split()
        for token in tokens:
            token = token.strip(" ,")
            # Process UART & USART
            if token == "UART*" or token == "USART*":
                found_names = GetBlockNamesStartingWith("USART*")
                found_names.extend(GetBlockNamesStartingWith("UART*"))
            else:
                found_names = GetBlockNamesStartingWith(token)
            # Check if found
            if len(found_names) == 0:
                print("No blocks found for {}".format(token))
                return
            needed_block.names.extend(found_names)
        # Check for duplicates
        if len(needed_block.names) != len(set(needed_block.names)):
            print("Duplicated block names: '{}' -> {}".format(needed_block.name_str, needed_block.names))
            return
        print("'{}' -> {}".format(needed_block.name_str, needed_block.names))

    # region ==== Process blocks-without-timers ====
    list_of_block_sets: list[list[ConfiguredBlock]] = []
    for needed_block in needed_blocks_without_timers:
        block_set: list[ConfiguredBlock] = []
        for avail_block in all_blocks:
            if avail_block.name in needed_block.names:
                # Available block found, add all its pin-signal-dma variants
                for pin_sig_dma in needed_block.pin_sig_dmas:
                    block_set.extend(avail_block.GetVariantsForPinsAndDmas(pin_sig_dma))
        if len(block_set) == 0:
            print("No variants found for {} with {}".format(needed_block.name_str, needed_block.pin_sig_dmas))
            return
        list_of_block_sets.append(block_set)

    for bset in list_of_block_sets:
        print("[", end="")
        for block in bset:
            print("{}_r{} ".format(block.name, block.remap_option), end="")
        print("]")

    combs: list[list[ConfiguredBlock]] = FindCombination(list_of_block_sets)
    # endregion

    # Find available blocks for timers
    for k, comb in enumerate(combs):
        # Build lists of used resources
        occupied_dmas = []
        occupied_pins = []
        for block in comb:
            occupied_dmas.extend(block.int_dmas)
            occupied_pins.extend(block.pinnums)
        # Find available blocks
        list_of_block_sets = []
        for tmr in needed_timers:
            usable_block_set = []  # list of usable blocks for the timer
            for avail_block in all_blocks:
                if avail_block.name in tmr.names:
                    # Available block found, check all its pin-signal-dma variants
                    for pin_sig_dma in tmr.pin_sig_dmas:
                        cfg_blocks = avail_block.GetVariantsForPinsAndDmas(pin_sig_dma)
                        for cblock in cfg_blocks:
                            # check dmas and pins
                            if not any(idma in occupied_dmas for idma in cblock.int_dmas):  # no duplicates
                                if not any(ipin in occupied_pins for ipin in cblock.pinnums):
                                    usable_block_set.append(cblock)
            if len(usable_block_set) == 0:
                print("No variants found for {} with {}".format(tmr.name_str, tmr.pin_sig_dmas))
                break  # Stop iterate timers and try next combination
            else:
                list_of_block_sets.append(usable_block_set)
        else:  # All timers have variants
            print("Combination {}:".format(k+1))
            for bset in list_of_block_sets:
                print("[", end="")
                for block in bset:
                    print("{}_r{} ".format(block.name, block.remap_option), end="")
                print("]")
            # Find non-crossing timers
            FindCombination(list_of_block_sets)


if __name__ == '__main__':
    main()
